import io
import os
import hmac
import math
from hashlib import pbkdf2_hmac, md5
from Crypto.Cipher import AES


PAGE_SIZE = 4096
SALT_SIZE = 16
IV_SIZE = 16
HMAC_SIZE = 64  # SHA512
RESERVE = IV_SIZE + HMAC_SIZE  # already multiple of 16 (80)
SQLITE_HEADER = b"SQLite format 3\x00"


def _derive_keys(key_bytes: bytes, salt: bytes) -> tuple[bytes, bytes]:
    enc_key = pbkdf2_hmac("sha512", key_bytes, salt, 256000, dklen=32)
    mac_salt = bytes([b ^ 0x3A for b in salt])
    mac_key = pbkdf2_hmac("sha512", enc_key, mac_salt, 2, dklen=32)
    return enc_key, mac_key


def _validate_key(first_page: bytes, key_bytes: bytes) -> bool:
    if len(first_page) < PAGE_SIZE or len(key_bytes) != 32:
        return False
    salt = first_page[:SALT_SIZE]
    _, mac_key = _derive_keys(key_bytes, salt)
    data_end = PAGE_SIZE - RESERVE + IV_SIZE
    mac = hmac.new(mac_key, first_page[SALT_SIZE:data_end], "sha512")
    mac.update((1).to_bytes(4, "little"))
    stored = first_page[data_end : data_end + HMAC_SIZE]
    return hmac.compare_digest(mac.digest(), stored)


def decrypt_database_file(input_path: str, hex_key: str, output_path: str) -> None:
    if len(hex_key) != 64:
        raise ValueError("密钥长度必须为64个十六进制字符")
    key_bytes = bytes.fromhex(hex_key)

    with open(input_path, "rb") as f_in, open(output_path, "wb") as f_out:
        first_page = f_in.read(PAGE_SIZE)
        if first_page.startswith(SQLITE_HEADER[:-1]):
            raise RuntimeError("数据库已经是解密状态")

        # 验证密钥
        if not _validate_key(first_page, key_bytes):
            raise RuntimeError("密钥验证失败")

        salt = first_page[:SALT_SIZE]
        enc_key, mac_key = _derive_keys(key_bytes, salt)

        # 写入SQLite头
        f_out.write(SQLITE_HEADER)

        # 回到文件开头重新逐页处理
        f_in.seek(0)
        page_index = 0
        while True:
            page = f_in.read(PAGE_SIZE)
            if not page:
                break

            if set(page) == {0}:
                f_out.write(page)
                page_index += 1
                continue

            offset = SALT_SIZE if page_index == 0 else 0
            data_end = PAGE_SIZE - RESERVE + IV_SIZE

            # 校验 HMAC
            mac = hmac.new(mac_key, page[offset:data_end], "sha512")
            mac.update((page_index + 1).to_bytes(4, "little"))
            if not hmac.compare_digest(mac.digest(), page[data_end : data_end + HMAC_SIZE]):
                raise RuntimeError("页面校验失败，HMAC 不匹配")

            iv = page[PAGE_SIZE - RESERVE : PAGE_SIZE - RESERVE + IV_SIZE]
            cipher = AES.new(enc_key, AES.MODE_CBC, iv=iv)
            encrypted = bytearray(page[offset : PAGE_SIZE - RESERVE])
            decrypted = cipher.decrypt(bytes(encrypted))

            # 拼回保留字节
            out_page = decrypted + page[PAGE_SIZE - RESERVE : PAGE_SIZE]
            f_out.write(out_page)
            page_index += 1


